package com.example.harmonet.harmtorch;

public class MaxPool1d implements Layer {

    int kernel_size;
    int stride;
    int padding;

    public MaxPool1d(int kernel_size, int stride, int padding) {
        this.kernel_size = kernel_size;
        this.stride = stride;
        this.padding = padding;
    }

    private float getMax(Tensor in, int n, int i, int j) throws Exception {
        float max = -Float.MAX_VALUE;
        for(int k = 0; k < kernel_size; ++k) {
            int pos = j * stride + k - padding;
            if(pos < 0) {
                max = Math.max(max, 0);
            }
            else if(pos < in.dim()[2]) {
                max = Math.max(max, in.tensor()[in.pos(n, i, pos)]);
            }
            else if(pos < in.dim()[2] + padding) {
                max = Math.max(max, 0);
            }
        }
        return max;
    }

    @Override
    public Tensor forward(Tensor in) throws Exception {
        int[] out_dim = new int[3];
        out_dim[0] = in.dim()[0];
        out_dim[1] = in.dim()[1];
        out_dim[2] = (in.dim()[2] + 2 * padding - kernel_size) / stride + 1;
        Tensor out = new Tensor(out_dim);
        for(int n = 0; n < out_dim[0]; ++n) {
            for(int i = 0; i < out_dim[1]; ++i) {
                for(int j = 0; j < out_dim[2]; ++j) {
                    out.tensor()[out.pos(n, i, j)] = getMax(in, n, i, j);
                }
            }
        }
        return out;
    }

    @Override
    public int getParam() {
        return 0;
    }

    @Override
    public void init(float[] param) {
    }
}
